#!/usr/bin/env python3
"""
Whisper 模型自动下载脚本

功能：
- 检测本地是否已有模型
- 自动使用国内镜像加速下载
- 支持下载所有模型大小
- 显示下载进度和模型位置
"""

import os
import sys
from pathlib import Path


def check_model_exists(model_size: str) -> bool:
    """检查模型是否已存在"""
    cache_dir = Path.home() / '.cache' / 'huggingface' / 'hub'

    # 查找模型目录
    if cache_dir.exists():
        for item in cache_dir.iterdir():
            if item.is_dir() and f'whisper-{model_size}' in item.name.lower():
                return True
    return False


def get_model_info(model_size: str) -> dict:
    """获取模型信息"""
    model_info = {
        'tiny': {'size': '~75 MB', 'speed': '最快', 'accuracy': '一般'},
        'base': {'size': '~145 MB', 'speed': '快', 'accuracy': '良好'},
        'small': {'size': '~485 MB', 'speed': '中等', 'accuracy': '很好'},
        'medium': {'size': '~1.5 GB', 'speed': '慢', 'accuracy': '优秀'},
        'large': {'size': '~3 GB', 'speed': '很慢', 'accuracy': '最佳'},
    }
    return model_info.get(model_size, {})


def download_model(model_size: str = 'tiny', use_mirror: bool = True):
    """
    下载 Whisper 模型

    Args:
        model_size: 模型大小 (tiny, base, small, medium, large)
        use_mirror: 是否使用国内镜像
    """
    print("=" * 60)
    print(f"  Whisper 模型下载器 - {model_size}")
    print("=" * 60)
    print()

    # 检查模型是否已存在
    if check_model_exists(model_size):
        print(f"✓ 检测到 {model_size} 模型已存在，无需重复下载")
        print()
        cache_dir = Path.home() / '.cache' / 'huggingface' / 'hub'
        print(f"模型位置: {cache_dir}")
        return True

    # 显示模型信息
    info = get_model_info(model_size)
    if info:
        print(f"模型大小: {info['size']}")
        print(f"转录速度: {info['speed']}")
        print(f"准确度: {info['accuracy']}")
        print()

    # 设置镜像源
    if use_mirror:
        print("✓ 使用国内镜像加速下载")
        os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
    else:
        print("⚠ 使用官方源下载（可能较慢）")

    print()
    print(f"开始下载 {model_size} 模型...")
    print("这可能需要几分钟，请耐心等待...")
    print("-" * 60)

    try:
        from faster_whisper import WhisperModel

        # 下载模型（会自动缓存）
        model = WhisperModel(
            model_size,
            device='cpu',
            compute_type='int8',
            download_root=None  # 使用默认缓存位置
        )

        print()
        print("=" * 60)
        print("  ✓ 模型下载成功！")
        print("=" * 60)
        print()
        print(f"模型位置: {Path.home() / '.cache' / 'huggingface'}")
        print()
        print("下一步：")
        print("  运行 Worker 即可自动使用该模型")
        print("  ./start_worker.sh --excel your_file.xlsx")
        print()

        return True

    except ImportError:
        print()
        print("❌ 错误: 未安装 faster-whisper")
        print()
        print("请先安装依赖:")
        print("  pip install faster-whisper")
        print()
        return False

    except Exception as e:
        print()
        print("=" * 60)
        print("  ❌ 模型下载失败")
        print("=" * 60)
        print()
        print(f"错误信息: {e}")
        print()
        print("可能的原因:")
        print("  1. 网络连接问题")
        print("  2. 磁盘空间不足")
        print("  3. 防火墙/代理限制")
        print()
        print("解决方案:")
        print("  1. 检查网络连接")
        print("  2. 如在企业网络，配置代理:")
        print("     set HTTP_PROXY=http://proxy:port")
        print("  3. 尝试使用镜像源 (默认已启用)")
        print()
        return False


def download_all_models(use_mirror: bool = True):
    """下载所有常用模型"""
    models = ['tiny', 'base', 'small']

    print("=" * 60)
    print("  批量下载 Whisper 模型")
    print("=" * 60)
    print()
    print(f"将下载以下模型: {', '.join(models)}")
    print()

    success = []
    failed = []

    for model in models:
        print(f"\n{'=' * 60}")
        print(f"  [{models.index(model) + 1}/{len(models)}] 下载 {model} 模型")
        print(f"{'=' * 60}\n")

        if download_model(model, use_mirror):
            success.append(model)
        else:
            failed.append(model)

    # 总结
    print("\n" + "=" * 60)
    print("  下载完成")
    print("=" * 60)
    print()
    print(f"成功: {len(success)} 个模型 - {', '.join(success)}")
    if failed:
        print(f"失败: {len(failed)} 个模型 - {', '.join(failed)}")
    print()


def main():
    """主函数"""
    import argparse

    parser = argparse.ArgumentParser(
        description='Whisper 模型自动下载器',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例:
  # 下载 tiny 模型（默认）
  python download_models.py

  # 下载 base 模型
  python download_models.py --model base

  # 下载所有常用模型
  python download_models.py --all

  # 不使用镜像（官方源）
  python download_models.py --no-mirror

可用模型:
  tiny   - 最快，约75MB，适合测试
  base   - 推荐，约145MB，准确度良好
  small  - 较慢，约485MB，准确度很好
  medium - 很慢，约1.5GB，准确度优秀
  large  - 最慢，约3GB，准确度最佳
        """
    )

    parser.add_argument(
        '--model', '-m',
        default='tiny',
        choices=['tiny', 'base', 'small', 'medium', 'large'],
        help='模型大小 (默认: tiny)'
    )

    parser.add_argument(
        '--all', '-a',
        action='store_true',
        help='下载所有常用模型 (tiny, base, small)'
    )

    parser.add_argument(
        '--no-mirror',
        action='store_true',
        help='不使用国内镜像，使用官方源'
    )

    parser.add_argument(
        '--check', '-c',
        action='store_true',
        help='仅检查模型是否存在，不下载'
    )

    args = parser.parse_args()

    use_mirror = not args.no_mirror

    # 仅检查模式
    if args.check:
        print("检查本地模型...")
        print()
        for model in ['tiny', 'base', 'small', 'medium', 'large']:
            exists = check_model_exists(model)
            status = "✓ 已安装" if exists else "✗ 未安装"
            print(f"{model:8s} - {status}")
        print()
        return

    # 下载模型
    if args.all:
        download_all_models(use_mirror)
    else:
        download_model(args.model, use_mirror)


if __name__ == '__main__':
    try:
        main()
    except KeyboardInterrupt:
        print("\n\n⚠ 已取消下载")
        sys.exit(1)
    except Exception as e:
        print(f"\n\n❌ 发生错误: {e}")
        sys.exit(1)
